__author__ = 'sdjkx'

import os
import shutil
from libs.ustr import *
from PyQt5.QtWidgets import QMessageBox,QFileDialog
try:
    import xml.etree.cElementTree as ET
except ImportError:
    import xml.etree.ElementTree as ET

class ScanPathToObjDetect:

    def __init__(self, parent ,label_dict):
        self.trainPath = "train"
        self.annotationsPath = "Annotations"
        self.JPEGImages = "JPEGImages"
        self.srcPath = ""
        self.targetPath = ""
        self.labelImg = parent
        self.label_dict = label_dict

    def getSrcPath(self):
        QMessageBox.information(self.labelImg, '提示', "请选择源图片路径和保存图片的目标路径，扫描源图片路径时将扫描指定路径的子目录")

        dirPath = ustr(QFileDialog.getExistingDirectory(self.labelImg,
                                                       "选择源图片路径", "",  QFileDialog.ShowDirsOnly
                                                       | QFileDialog.DontResolveSymlinks))

        if dirPath is None or len(dirPath) <= 1:
            QMessageBox.information(self.labelImg, '错误', "必须指定源图片路径！")
            return
        else:
            self.srcPath = dirPath

        dirPath = ustr(QFileDialog.getExistingDirectory(self.labelImg,
                                                       "选择目标路径", "",  QFileDialog.ShowDirsOnly
                                                       | QFileDialog.DontResolveSymlinks))

        if dirPath is None or len(dirPath) <= 1:
            QMessageBox.information(self.labelImg, '错误', "必须指定目标路径！")
            return
        else:
            self.targetPath = dirPath

        ret = QMessageBox.question(self.labelImg, '提示', "确定要扫描目录【"+self.srcPath+"】，把图片保存到【"+self.targetPath+"】吗？",
                                   QMessageBox.Yes | QMessageBox.No, QMessageBox.No)
        if ret == QMessageBox.Yes:
            self.scan()
            return

    def scan(self):
        self.trainPathList = []
        self.scanDir(self.srcPath)

        for path in self.trainPathList:
            self.copyJpgAndXMl(path)

        self.createTrainValTest()
        self.createLabel()

        QMessageBox.information(self.labelImg, '提示', "数据生成完成！")

    def scanDir(self, srcPath):
        path = os.walk(srcPath)

        for root, dirs, files in path:
            if len(dirs) > 0:
                for dir in dirs:
                    if dir == self.trainPath:
                        newPath = os.path.join(root, dir)
                        self.trainPathList.append(newPath)
                    else:
                        newPath = os.path.join(root, dir)
                        self.scanDir(newPath)

    def copyJpgAndXMl(self, root):
        print("复制目录文件：" + root)
        newJpgPath = root # os.path.join(root, self.trainPath)
        if not os.path.exists(newJpgPath):
            QMessageBox.information(self.labelImg, '错误',"路径【"+newJpgPath+"】不存在 ")
            return

        xmlPath =  newJpgPath[:len(newJpgPath) - len(self.trainPath)] + self.annotationsPath
        if not os.path.exists(xmlPath):
            QMessageBox.information(self.labelImg, '错误', "路径【"+root+"】没有相应的【"+self.annotationsPath+"】目录！")
            return

        jpgTargetPath = os.path.join(self.targetPath, self.JPEGImages)
        xmlTargetPath = os.path.join(self.targetPath, self.annotationsPath)
        # 如果不存在目的目录则创建一个，保持层级结构
        if not os.path.exists(jpgTargetPath):
            os.makedirs(jpgTargetPath)
        if not os.path.exists(xmlTargetPath):
            os.makedirs(xmlTargetPath)

        print("开始复制文件路径：" + newJpgPath + "        XMLPath = " + xmlPath)
        for filename in os.listdir(newJpgPath):
            srcJpgFile = os.path.join(newJpgPath, filename)
            targetJpgFile = os.path.join(jpgTargetPath, filename)

            xmlFileName = filename[:len(filename) - 3] + "xml"
            srcXmlFile = os.path.join(xmlPath, xmlFileName)
            xmlTargetFile = os.path.join(xmlTargetPath, xmlFileName)

            if os.path.exists(srcJpgFile) and os.path.exists(srcXmlFile):
                if not os.path.exists(targetJpgFile):
                    shutil.copy(srcJpgFile, targetJpgFile)
                self.readXmlToTarget(srcXmlFile, xmlTargetFile)

    def readXmlToTarget(self, srcXmlFile, xmlTargetFile):
        # print("开始复制XML文件：【" + srcXmlFile + "】")
        tree = ET.parse(srcXmlFile)             # 打开 xml 文档
        root = tree.getroot()                   # 获得 root 节点

        for object in root.findall('object'):   # 找到root节点下的所有object节点
            name = object.find('name').text     # 子节点下节点name的值

            if name in self.label_dict:
                label =  self.label_dict[name]
                enName = label.get_labelInit()
                label.labelCount += 1
            else:
                print("文件【" + srcXmlFile + "】中的标签名称【" + name + "】没有对应的字典数据")
                return

            object.find('name').text = enName

        # pass
        tree.write(xmlTargetFile)          # 保存文件

    def createTrainValTest(self):
        print("开始生成图片集、测试集数据文件")
        jpgPath = os.path.join(self.targetPath, self.JPEGImages)
        # xmlPath = os.path.join(self.targetPath, self.annotationsPath)

        all_tran_str = ""
        all_val_str = ""
        all_test_str = ""
        i = 0
        for filename in os.listdir(jpgPath):
            i += 1

            xmlFileName = filename[:len(filename) - 3] + "xml"
            str = "./JPEGImages/" + filename+'\t./Annotations/' + xmlFileName + "\n"
            all_tran_str += str
            if i % 50 == 0:
                all_test_str += str
            if i % 20 == 0:
                all_val_str += str

        out_txt_path_train = os.path.join(self.targetPath, 'train.txt')
        out_txt_path_test = os.path.join(self.targetPath,'test.txt')
        out_txt_path_val = os.path.join(self.targetPath,'val.txt')

        txt_train = open(out_txt_path_train, 'w')
        txt_test = open(out_txt_path_test, 'w')
        txt_val = open(out_txt_path_val, 'w')
        txt_train.write(all_tran_str)
        txt_test.write(all_test_str)
        txt_val.write(all_val_str)
        txt_train.close()
        txt_test.close()
        txt_val.close()
        print("生成图片集、测试集数据文件完成")

    def createLabel(self):
        print("开始生成Label文件")
        label_str = ""

        for label in self.label_dict.values():
            if label.labelCount > 0:
                label_str += label.get_labelInit() + "\n"

        out_txt_path_label = os.path.join(self.targetPath, 'label_list.txt')
        txt_train = open(out_txt_path_label, 'w')
        txt_train.write(label_str)
        txt_train.close()
        print("生成Label文件----完成")
